from xml.etree import ElementTree as ET
from pymongo import MongoClient

mongo_server='10.1.40.200'
ip_address='10.1.40.81'
apikey="2d5ec79df8a111ea97be005056b3e852"

port="7508"
prod_id=666
col_id=1442
test_protocol="http"


def pretty_xml(element, indent, newline, level=0):  # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行
    if element:  # 判断element是否有子元素
        if (element.text is None) or element.text.isspace():  # 如果element的text没有内容
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
            # else:  # 此处两行如果把注释去掉，Element的text也会另起一行
            # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element)  # 将element转成list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1):  # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
            subelement.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            subelement.tail = newline + indent * level
        pretty_xml(subelement, indent, newline, level=level + 1)  # 对子元素进行递归操作


def create_env_element(root,param_name,param_value):
    elementProp = ET.SubElement(root, 'elementProp', attrib={"name": param_name, "elementType": "Argument"})
    stringProp = ET.SubElement(elementProp, 'stringProp', attrib={"name": "Argument.name"})
    stringProp.text = param_name
    stringProp = ET.SubElement(elementProp, 'stringProp', attrib={"name": "Argument.value"})
    stringProp.text = param_value
    stringProp = ET.SubElement(elementProp, 'stringProp', attrib={"name": "Argument.metadata"})
    stringProp.text = '='


def create_string_param(root,tag_name="stringProp",value=None,attrib=None):
    #封闭式数据，带有值和tag
    if attrib is None:
        stringProp = ET.SubElement(root, tag_name)
    else:
        stringProp = ET.SubElement(root, tag_name, attrib=attrib)
    if value:
        stringProp.text = value
    return stringProp


def create_loop_element(root,suite_name,param_value=None):
    ThreadGroup = ET.SubElement(root, 'ThreadGroup', attrib={"testname": suite_name, "guiclass":"ThreadGroupGui",
                                                             "testclass":"ThreadGroup" ,"enabled":"true"} )

    create_string_param(ThreadGroup,attrib={"name":"ThreadGroup.on_sample_error"},value="continue")
    elementProp=ET.SubElement(ThreadGroup,'elementProp',attrib={"name":"ThreadGroup.main_controller","elementType":"LoopController",
                                                                "guiclass":"LoopControlPanel","testclass":"LoopController",
                                                                "testname":"循环控制器","enabled":"true"})
    create_string_param(elementProp, attrib={"name":"LoopController.continue_forever"}, value="false", tag_name="boolProp")
    create_string_param(elementProp, attrib={"name":"LoopController.loops"}, value="1")
    create_string_param(ThreadGroup, attrib={"name":"ThreadGroup.num_threads"}, value="1")
    create_string_param(ThreadGroup, attrib={"name":"ThreadGroup.ramp_time"}, value="1")
    create_string_param(ThreadGroup, attrib={"name":"ThreadGroup.scheduler"}, value="false")
    create_string_param(ThreadGroup, attrib={"name":"ThreadGroup.duration"}, value="20")
    create_string_param(ThreadGroup, attrib={"name":"ThreadGroup.delay"})


def create_header_element(root,suite_name,headers,parameters,body,interface_id,ip_address, port, test_protocol,db,comment="",):
    table = db["interface"]
    x=table.find_one({'_id': int(interface_id)},{'method':1,'path':1})
    print(x)
    method=x['method']
    path=x['path']

    hashTree = ET.SubElement(root, 'hashTree')

    #头信息树
    HeaderManager = ET.SubElement(hashTree, 'HeaderManager', attrib={"guiclass":"HeaderPanel","testclass":"HeaderManager",
                                                                     "testname":"HTTP Header Manager","enabled":"true"})
    collectionProp=ET.SubElement(HeaderManager,'collectionProp',attrib={"name":"HeaderManager.headers"})

    #自定义的头信息
    for k,v in headers.items():
        elementProp = create_string_param(collectionProp, tag_name="elementProp",
                                      attrib={"name": "", "elementType": "Header"})
        create_string_param(elementProp, attrib={"name": "Header.value"}, value=v)
        create_string_param(elementProp, attrib={"name": "Header.name"}, value=k)

    #请求信息树
    hashTree = ET.SubElement(hashTree, 'hashTree')
    HTTPSamplerProxy = ET.SubElement(hashTree, 'HTTPSamplerProxy', attrib={"guiclass": "HttpTestSampleGui",
                                                                           "testclass": "HTTPSamplerProxy",
                                                                           "testname": suite_name,
                                                                           "enabled": "true"})

    elementProp=ET.SubElement(HTTPSamplerProxy, 'elementProp',
                              attrib={"name": "HTTPsampler.Arguments","elementType" :"Arguments",
                                      "guiclass":"HTTPArgumentsPanel","testclass":"Arguments","testname":"User Defined Variables","enabled":"true"})
    collectionProp = ET.SubElement(elementProp, 'collectionProp',attrib={"name": "Arguments.arguments"})
    #自定义的query参数
    if body in ({},None,""):
        for k,v in parameters.items():
            min_elementProp = ET.SubElement(collectionProp, 'elementProp',
                                       attrib={"name": k,"elementType":"HTTPArgument"})
            create_string_param(min_elementProp, attrib={"name": "HTTPArgument.always_encode"}, value="false",
                                tag_name="boolProp")
            create_string_param(min_elementProp, attrib={"name": "Argument.value"}, value=v, tag_name="stringProp")
            create_string_param(min_elementProp, attrib={"name": "Argument.metadata"}, value="=", tag_name="stringProp")
            create_string_param(min_elementProp, attrib={"name": "HTTPArgument.use_equals"}, value="true", tag_name="boolProp")
            create_string_param(min_elementProp, attrib={"name": "Argument.name"}, value=k, tag_name="stringProp")
    else:
        #自定义的body参数
        min_elementProp = ET.SubElement(collectionProp, 'elementProp',
                                            attrib={"name": "", "elementType": "HTTPArgument"})
        create_string_param(min_elementProp, attrib={"name": "HTTPArgument.always_encode"}, value="false",
                                tag_name="boolProp")
        create_string_param(min_elementProp, attrib={"name": "Argument.value"}, value=body, tag_name="stringProp")
        create_string_param(min_elementProp, attrib={"name": "Argument.metadata"}, value="=", tag_name="stringProp")
        create_string_param(min_elementProp, attrib={"name": "HTTPArgument.use_equals"}, value="true",
                                tag_name="boolProp")


    #基础请求信息
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.domain"}, value=ip_address, tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.port"}, value=port, tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.protocol"}, value=test_protocol, tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.contentEncoding"}, value="utf-8", tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.path"}, value=path,tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.method"}, value=method, tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.follow_redirects"}, value="false", tag_name="boolProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.auto_redirects"}, value="true",
                        tag_name="boolProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.use_keepalive"}, value="true",
                        tag_name="boolProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.DO_MULTIPART_POST"}, value="false",
                        tag_name="boolProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.embedded_url_re"}, value="",
                        tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.implementation"}, value="Java",
                        tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.connect_timeout"}, value="",
                        tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "HTTPSampler.response_timeout"}, value="",
                        tag_name="stringProp")
    create_string_param(HTTPSamplerProxy, attrib={"name": "TestPlan.comments"}, value=comment,
                        tag_name="stringProp")

def create_suit(root,db,col_id,prod_id,ip_address, port, test_protocol):
    #sql="db.getCollection('interface_case').find({'project_id':{'$in':[666]}},{'case_env':1,'project_id':1,'interface_id':1,'casename':1,'req_body_other':1,'req_params':1,'req_headers':1,'req_query':1,'req_body_form':1,'test_res_body':1,'test_res_header':1})"
    table = db["interface_case"]
    x = table.find({'$and':[{'project_id':{'$in':[int(prod_id)]}},{'col_id':{'$in':[int(col_id)]}}]},{'case_env':1,'project_id':1,'interface_id':1,'casename':1,'req_body_other':1,'req_params':1,'req_headers':1,'req_query':1,'req_body_form':1,'test_res_body':1,'test_res_header':1})

    order=0
    for i in x:
        order+=1
        print(i)
        headers={}
        parameters={}
        for j in i['req_headers']:
            headers.update({j['name']:j.get('value',"")})
        for j in i['req_query']:
            parameters.update({j['name']:j.get('value',"")})
        body=i.get('req_body_other',None)
        suite_name='测试用例{}-{}'.format(order,i['casename'] )
        create_loop_element(root, suite_name)
        create_header_element(root, "发起HTTP请求",headers,parameters,body,i['interface_id'],ip_address, port, test_protocol,db)


def create_result_view(root):
    ResultCollector = ET.SubElement(root, 'ResultCollector',
                                  attrib={"guiclass":"ViewResultsFullVisualizer","testclass":"ResultCollector",
                                          "testname":"View Results Tree","enabled":"true"})
    create_string_param(ResultCollector, attrib={"name":"ResultCollector.error_logging"}, value="false",tag_name="boolProp")
    objProp=ET.SubElement(ResultCollector, 'objProp')
    name = ET.SubElement(objProp, 'name')
    name.text="saveConfig"
    value=ET.SubElement(objProp,'value',attrib={"class":"SampleSaveConfiguration"})
    tags=['time','latency','timestamp','success','label','code','message','threadName','dataType',
          'assertions','subresults','fieldNames','saveAssertionResultsFailureMessage','sentBytes',
          'bytes','url','threadCounts','idleTime','connectTime']
    for i in tags:
        temp=ET.SubElement(value, i)
        temp.text="true"
    tags=['encoding','responseData','samplerData','xml','responseHeaders','requestHeaders','responseDataOnError']
    for i in tags:
        temp=ET.SubElement(value, i)
        temp.text = "false"
    temp=ET.SubElement(value, 'assertionsResultsToSave')
    temp.text = "0"
    create_string_param(ResultCollector, attrib={"name": "filename"}, tag_name="stringProp")
    hashTree = ET.SubElement(root, 'hashTree')


def writeXML(prod_id,col_id,ip_address,apikey,port,test_protocol):
    global mongo_server
    client = MongoClient("mongodb://{}:27017".format(mongo_server))
    dblist = client.list_database_names()
    print(dblist)
    db = client["yapi"]
    table=db['project']
    x=table.find_one({'_id': prod_id}, {'name': 1})
    prod_name=x['name']

    file_name='{}测试脚本.jmx'.format(prod_name)

    root = ET.Element('jmeterTestPlan',attrib={"version":"1.2","properties":"5.0","jmeter":"5.1.1 r1855137"})  # 创建首节点
    main_hashTree = ET.SubElement(root, 'hashTree')  # 增加子节点

    item = ET.SubElement(main_hashTree, 'TestPlan', attrib={"guiclass":"TestPlanGui", "testclass":"TestPlan",
                                                       "testname":"Test产品基准测试", "enabled":"true"})
    stringProp = ET.SubElement(item, 'stringProp', attrib={"name":"TestPlan.comments"})
    boolProp=ET.SubElement(item, 'boolProp', attrib={"name":"TestPlan.functional_mode"})
    boolProp.text='false'
    boolProp = ET.SubElement(item, 'boolProp', attrib={"name": "TestPlan.tearDown_on_shutdown"})
    boolProp.text = 'true'
    boolProp = ET.SubElement(item, 'boolProp', attrib={"name": "TestPlan.serialize_threadgroups"})
    boolProp.text = 'false'
    elementProp=ET.SubElement(item, 'elementProp', attrib={"name":"TestPlan.user_defined_variables","elementType":"Arguments",
                                                           "guiclass":"ArgumentsPanel","testclass":"Arguments",
                                                           "testname":"User Defined Variables","enabled":"true"})
    collectionProp=ET.SubElement(elementProp, 'collectionProp',attrib={"name":"Arguments.arguments"})
    create_env_element(collectionProp,'http',test_protocol)
    create_env_element(collectionProp, 'ip', ip_address)
    create_env_element(collectionProp, 'apikey', apikey)
    create_env_element(collectionProp, 'port', port)
    stringProp=ET.SubElement(item, 'stringProp', attrib={"name":"TestPlan.user_define_classpath"})

    hashTree = ET.SubElement(main_hashTree, 'hashTree')

    #加入测试用例
    create_suit(hashTree,db,col_id,prod_id,ip_address, port, test_protocol)
    #加入结果展示器
    create_result_view(hashTree)




    # create_suite_element(hashTree, "testname")

    # item_3 = ET.Element('item', {'id': '2'})  # 创建节点
    # root.append(item_3)  # 追加节点到root节点下
    # caption_2 = ET.Element('caption')
    # caption_2.text = 'Zope'
    # item_3.append(caption_2)

    w = ET.ElementTree(root)
    w.write(file_name, 'utf-8', True)
    # 代码美化
    tree = ET.parse(file_name)  # 解析movies.xml这个文件
    root = tree.getroot()  # 得到根元素，Element类
    pretty_xml(root, '\t', '\n')  # 执行美化方法
    tree.write(file_name)
    return file_name


if __name__ == '__main__':
    writeXML(prod_id,col_id,ip_address,apikey,port,test_protocol)



